package com.lw.leetcode.sort.b;

import java.util.Map;
import java.util.TreeMap;

/**
 * Created with IntelliJ IDEA.
 * 6113. 无限集中的最小数字
 *
 * @author liw
 * @version 1.0
 * @date 2022/7/11 9:09
 */
public class SmallestInfiniteSet {


    public static void main(String[] args) {
        SmallestInfiniteSet test = new SmallestInfiniteSet();

        // [null, null, 1, 2, 3, null, 1, 4, 5]
        test.addBack(2);
        System.out.println(test.popSmallest());
        System.out.println(test.map);
        System.out.println(test.popSmallest());
        System.out.println(test.map);
        System.out.println(test.popSmallest());
        System.out.println(test.map);
        test.addBack(1);
        System.out.println(test.map);
        System.out.println(test.popSmallest());
        System.out.println(test.map);
        System.out.println(test.popSmallest());
        System.out.println(test.map);
        System.out.println(test.popSmallest());
        System.out.println(test.map);

    }

    private TreeMap<Integer, Integer> map = new TreeMap<>();

    public SmallestInfiniteSet() {

    }

    public int popSmallest() {
        Map.Entry<Integer, Integer> entry = map.ceilingEntry(0);
        if (entry == null) {
            map.put(1, 1);
            return 1;
        }
        Integer key = entry.getKey();
        if (key != 1) {
            if (key == 2) {
                map.put(1, map.get(2));
                map.remove(2);
            } else {
                map.put(1, 1);
            }
            return 1;
        }
        Integer value = entry.getValue();
        entry = map.ceilingEntry(value + 1);
        if (entry == null || entry.getKey() > value + 2) {
            map.put(1, value + 1);
            return value + 1;
        }
        map.remove(entry.getKey());
        map.put(1, entry.getValue());
        return value + 1;
    }

    public void addBack(int num) {
        Map.Entry<Integer, Integer> entry = map.floorEntry(num);
        if (entry != null) {
            int key = entry.getKey();
            int value = entry.getValue();
            if (value >= num) {
                if (key == value) {
                    map.remove(key);
                    return;
                }
                if (num == key) {
                    map.remove(key);
                    map.put(key + 1, value);
                    return;
                }
                if (num == value) {
                    map.put(key, value - 1);
                    return;
                }
                map.put(key, num - 1);
                map.put(num + 1, value);
            }
        }
    }

}
